from ase.io import read, write
import numpy as np

atoms_list = read("train_small.xyz", ":")
for atoms in atoms_list:
    if "forces" in atoms.arrays:
        atoms.arrays["forces"] = atoms.arrays["forces"].astype(np.float32)
    if "energy" in atoms.info:
        atoms.info["energy"] = np.float32(atoms.info["energy"])
write("train_small_fp32.xyz", atoms_list)
